use "../Common/CG.sml";

(* structure CoreEx is used for extending the Core language AST with a context-graph "attached" to each node.
   It is used in the abstract-interpreter algorithm to hold the partial / temporary and eventually final interpretation 
   result graphs of each of the AST nodes analysis. *)
structure CoreEx = 
struct
	(* datatype used for defining the AST extended with context-graph "attached" to each node. *)
	datatype CmdEx = Skip | Asgn of int * Core.Expr | Choice of Cmd * Cmd |
	       Seq  of Cmd * Cmd  | Loop of int * Cmd
	and
		   (* Constructor for a command with a graph "attached". *)
	       Cmd = cmd of CmdEx * CG.graph ref;
	
	(* converts a regular AST to an extended AST.
       Note:  loop body command c is expanded to: "if { c } else { Skip }".
	   This is done to avoid calculating the Skip graph explicitly by the abstact interpreter algorithm. *)
	fun extendedAST(Core.Skip) = cmd(Skip, ref (CG.newGraph([])))
	  | extendedAST(Core.Asgn(i,e)) = cmd(Asgn(i,e), ref (CG.newGraph([])))
	  | extendedAST(Core.Choice(c1,c2)) = cmd(Choice(extendedAST(c1), extendedAST(c2)), ref (CG.newGraph([])))
	  | extendedAST(Core.Seq(c1,c2)) = cmd(Seq(extendedAST(c1), extendedAST(c2)), ref (CG.newGraph([])))
	  | extendedAST(Core.Loop(i,c)) = cmd(Loop(i, extendedAST(Core.Choice(Core.Skip,c))), ref (CG.newGraph([]))); 
	
	(* converts an extended AST to a regular AST. This function is not the inverse of the extendedAST function (due to the loop expansion). *)
	fun reducedAST(cmd(Skip, _)) = Core.Skip
	  | reducedAST(cmd(Asgn(i,e),_)) = Core.Asgn(i,e)
	  | reducedAST(cmd(Choice(c1,c2),_)) = Core.Choice(reducedAST(c1),reducedAST(c2))
	  | reducedAST(cmd(Seq(c1,c2),_)) = Core.Seq(reducedAST(c1),reducedAST(c2))
	  | reducedAST(cmd(Loop(i,c),_)) = Core.Loop(i,reducedAST(c));

	(* extracts the graph of a CoreEx.Cmd *)
	fun graph(cmd(c,g)) = g
	
	(* prints the reduced AST of c *)
	fun prt(c) = Core.prt(reducedAST(c));  
end;

(* composition operator for context-graphs *)
infix o;
fun ((g1 : CG.graph) o (g2 : CG.graph)) : CG.graph = CG.compose(g1,g2);

(* LUB (join) operator for context-graphs *)
infix ||;
fun ((g1 : CG.graph) || (g2 : CG.graph)) : CG.graph = CG.join(g1,g2);

(* structure AbstractInterpreter encapsulates the abstract interpreter algorithm, as described in abstractInterpreterAlgorithm.docx *)
structure AbstractInterpreter =
struct
structure analysisKind = 
	struct 
		val nonlinear_Dvalue = D.Two;
		val name = "POLYNOMIAL";
	end; 

(* a global identifier which is updated (at the beginning of the algorithm execution) 
   to hold the number of variables in the analysed program. *)
val numVars = ref 0;
(* an iterator for debug print-out which is increased on each analysis interation over a loop body. *)
val itCount = ref 0;
(* debug mode -  when true - each algorithm stage will be printed, along with its output (and sometimes input). *)
val debug = CG.debug;	

(* prints s to the stdOut if debug is true. *)
fun debugPrint(s) = if (!debug) then print(s) else ();

(* returns flow i-(1)->i if i is not in source context and i <> x *)
fun computeFlowFromItoI(i : int, x : int, srcContext : CG.Context) =
	if (i <> x andalso (not (CG.inContext(i,srcContext)))) then [(i,i,D.One)] else [];

(* computes context graph for commands of the form "X := 0" *)
fun computeAsgnZeroGraph(x : int, preContexts : CG.Context list) : CG.graph =
	let 
		fun computeDataFlows(srcContext : CG.Context) : (int * int * D.t) list =
			let
				fun computeMultipleFlows(0,_) = []
                  |	computeMultipleFlows(i : int, srcContext : CG.Context) =
						computeFlowFromItoI(i,x,srcContext) @ computeMultipleFlows(i-1,srcContext);
			in
				computeMultipleFlows((!numVars),srcContext)
			end;
			
		fun computeEdge(srcContext : CG.Context) = 
			let
			    (* create target context identical to source context, at least at first... *)
				val tgtContext = CG.copyContext(!numVars, srcContext);
				(* include x in target context (equivalent to assignment of X := 0) *)
				val _ = CG.addVarToContext(x,tgtContext);
				val m = mat.M(computeDataFlows(srcContext));
			in
				CG.E ( srcContext, tgtContext, ref (sop_bar.dfr(m, mat.r(m))) )
			end;
	in
		CG.newGraph(map computeEdge preContexts)
	end;	

(* computes context graph for commands of the form "X := Y" *)
fun computeAsgnVarGraph(x : int, y : int, preContexts : CG.Context list) : CG.graph =
	let 
		fun computeDataFlows(srcContext : CG.Context) : (int * int * D.t) list =
			let
				fun computeFlowFromYtoX(i : int, srcContext : CG.Context)  =
					if (i = y andalso (not (CG.inContext(y, srcContext)))) then 
						[(y,x,D.One)]
					else 
						[];

				
				fun computeMultipleFlows(0,_) = []
                  |	computeMultipleFlows(i : int, srcContext : CG.Context) =
						let
							val flowYtoX = computeFlowFromYtoX(i,srcContext);
							val flowItoI = computeFlowFromItoI(i,x,srcContext);
						in
							flowYtoX @ flowItoI @ computeMultipleFlows(i-1, srcContext)
						end;
			in
				computeMultipleFlows((!numVars), srcContext)
			end;
			
		fun computeEdge(srcContext : CG.Context) = 
			let
			    (* create target context identical to source context, at least at first... *)
				val tgtContext = CG.copyContext(!numVars, srcContext);
				val y_in_srcContext = CG.inContext(y, srcContext);
				(* add x to target context if y is in source context (equivalent to assignment of X := 0),
				    otherwise exclude x from target context (equivalent to assignment of X := Y, where Y <> 0) *)
				val () = CG.setVarInContext(x, tgtContext, y_in_srcContext); 
				val m = mat.M(computeDataFlows(srcContext));
			in
				CG.E ( srcContext, tgtContext, ref (sop_bar.dfr(m, [])) )
			end;
	in
		CG.newGraph(map computeEdge preContexts)
	end;	

(* computes context graph for commands of the form "X := Y * Z" *)
fun computeAsgnTimesGraph(x, y, z, preContexts) : CG.graph =
	let 
		fun computeDataFlows(srcContext : CG.Context) : (int * int * D.t) list =
			let
				fun computeSingleFlow(i : int, srcContext : CG.Context) : (int * int * D.t) list =
					let 
						(* add flow y-(2)->x (or z-(2)->x) if y and z are both not in the source context (equivalent to assignment of X := Y*Z where Y>0 and Z>0 ) *)
						fun computeFlowFromYandZtoX(i : int, srcContext : CG.Context) =
							if ( (i = y orelse i = z) andalso
							     (not (CG.inContext(y, srcContext) orelse CG.inContext(z, srcContext))) )
							then [(i,x,D.Two)]
							else [];
					in
						computeFlowFromYandZtoX(i,srcContext) @ computeFlowFromItoI(i,x,srcContext)
					end;
				
				fun computeMultipleFlows(0,_) = []
                  |	computeMultipleFlows(i : int, srcContext : CG.Context) =
						computeSingleFlow(i, srcContext) @ computeMultipleFlows(i-1, srcContext);
			in
				computeMultipleFlows((!numVars), srcContext)
			end;
			
		fun computeEdge(srcContext : CG.Context) = 
			let
			    (* create target context identical to source context, at least at first... *)
				val tgtContext = CG.copyContext(!numVars, srcContext);
				val y_or_z_in_srcContext = (CG.inContext(y, srcContext)) orelse (CG.inContext(z, srcContext));
				(* add x to target context if  y or z are in the source context (equivalent to assignment of X := 0),
				    otherwise exclude x from target context (equivalent to assignment of X := Y * Z, where Y>0 and Z > 0) *)
				val _ = CG.setVarInContext(x, tgtContext, y_or_z_in_srcContext); 
				val mat = mat.M(computeDataFlows(srcContext));
			in
				CG.E ( srcContext, tgtContext, ref (sop_bar.dfr(mat, [])) )
			end;
	in
		CG.newGraph( map computeEdge preContexts )
	end;	

(* computes context graph for commands of the form "X := Y + Z" *)	
fun computeAsgnPlusGraph(x, y, z, preContexts) : CG.graph =
	let 
		fun computeDataFlows(x : int, y : int, z : int, srcContext : CG.Context) : (int * int * D.t) list =
			let
				fun computeSingleFlow(i : int, srcContext : CG.Context) : (int * int * D.t) list =
					let 
						(* add flow y-(1')->x (or z-(1')->x) if y (or z) are not in the source context (equivalent to assignment of X := Y+Z where Y>0 or Z>0 ) *)
						fun computeFlowFromYandZtoX(i : int, srcContext : CG.Context) =
							let
								val iNotInSrcContext = not (CG.inContext(i,srcContext));
								val yInSrcContext = CG.inContext(y, srcContext);
								val zInSrcContext = CG.inContext(z, srcContext);
							in
								if (i = y orelse i = z) then
									if (iNotInSrcContext andalso (not (yInSrcContext andalso zInSrcContext))) then 
										(if ((not yInSrcContext) andalso (not zInSrcContext)) then
											[(i,x,D.One')]
										else 
											[(i,x,D.One)])
									else []
								else []
							end;
					in
						 computeFlowFromYandZtoX(i,srcContext) @ computeFlowFromItoI(i,x,srcContext)
					end;
				
				fun computeMultipleFlows(0,_) = []
                  |	computeMultipleFlows(i : int, srcContext : CG.Context) =
							computeSingleFlow(i,srcContext) @ computeMultipleFlows(i-1,srcContext);								  

			in
				computeMultipleFlows((!numVars),srcContext)
			end;
			
		fun computeEdge(srcContext : CG.Context) = 
			let
			    (* create target context identical to source context, at least at first... *)
				val tgtContext = CG.copyContext(!numVars, srcContext);
				val y_and_z_in_srcContext = (CG.inContext(y, srcContext)) andalso (CG.inContext(z, srcContext));
				(* add x to target context if  y and z are both in the source context (equivalent to assignment of X := 0),
				    otherwise exclude x from target context (equivalent to assignment of X := Y + Z, where Y>0 or Z > 0) *)
				val _ = CG.setVarInContext(x, tgtContext, y_and_z_in_srcContext);  
				val mat = mat.M(computeDataFlows(x,y,z,srcContext));
			in
				CG.E ( srcContext, tgtContext, ref (sop_bar.dfr(mat, mat.r(mat))) )
			end;
			
	in
		CG.newGraph( map computeEdge preContexts )
	end;

(* updates DFRs on all edges of graph g with all possible data flows of the form (i->j,i'->j'). *)
fun computeDoubleFlows(g : CG.graph) =
	let
		fun updateDoubleFlowsOnEdge(CG.E (p, q, dfr as (ref (sop_bar.dfr(m,_)))) ) =
			let	
				(* filter function to determine if a given double data flow should be included in the result *)
				fun includeFlow (p : CG.Context,q : CG.Context) ((i,j),(i',j')) = 
								(not (CG.inContext(i,p))) andalso (not (CG.inContext(i',p))) andalso
								(not (CG.inContext(j,q))) andalso (not (CG.inContext(j',q))) andalso
								(i <> i' orelse j <> j');

				(* filter double flows according to source and target contexts *)
				val r' = List.filter (includeFlow(p,q)) (mat.r(m)); 				
			in
				dfr := sop_bar.dfr(m,r')		
			end;
	in
		List.app updateDoubleFlowsOnEdge (CG.edges(g))
	end;

(* computes a graph corresponding to the abstract interpretation of an assignment command, with given pre-contexts *)
fun computeAssignmentGraph (x : int, e : Core.Expr, preContexts : CG.Context list) : CG.graph ref =
	let
		(* computes data flows of the form i-(d)->j *)
		fun computeSingleFlows (x : int, e : Core.Expr, preContexts : CG.Context list) : CG.graph =
			case e of Core.Zero =>  						computeAsgnZeroGraph(x, preContexts)
					| Core.Var y =>  						computeAsgnVarGraph(x, y, preContexts)
					| Core.times(Core.Var y, Core.Var z) => computeAsgnTimesGraph(x, y, z, preContexts)									  
					| Core.plus(Core.Var y , Core.Var z) => computeAsgnPlusGraph(x, y, z, preContexts);
		
		val g' = ref (computeSingleFlows(x, e, preContexts));
		(* computes all data flows of the form (i->j,i'->j') over the flows calculated for all edges of the graph g' *)
		val () = computeDoubleFlows(!g');
	in
	   g'
	end;          

(* abstractly interpret command c,  creating (or updating) its CG.graph, return ref to resulting graph *)
fun abstractInterpret(c : CoreEx.Cmd, preContexts : CG.Context list) : CG.graph ref = 
	let
		(* Computes the loop correction on all edges where applicable. This function applies loop rule L2 in the DICE paper. *)
		fun doLoopCorrection(l : int, g : CG.graph ref)  =
			let
				(* override operator * to represent DFR composition, in thie scope of this function only *)
				fun ((d1 : sop_bar.DFR) * (d2 : sop_bar.DFR)) : sop_bar.DFR = sop_bar.prod(d1,d2)
				
				val sortedEdges = CG.sortEdges(CG.edges(!g));
				
				(* The loop correction operator on DFRs. *)
				fun LC(l,sop_bar.dfr(m,r)) =
					let
						val m' = mat.LC (l, analysisKind.nonlinear_Dvalue) m;
					in
						sop_bar.dfr(m', sop_bar.filterR(r,m'))
					end;
				
				(* computes LC on all edges (p,q') s.t. (p,q), (q,q') and (q,q) exist on graph *)
				fun computeSingleEdge  [] _ = []
				  | computeSingleEdge ((CG.E(p',q',d'))::t) (e as (CG.E (p,q,d))) = 
						(if (CG.eqContexts(q,p')) then
							(case CG.getEdgeTag(!g,CG.UE (q,q)) of
											SOME dfr => [(CG.E (p,q', ref ( !d * LC(l, !dfr) * !d' )))] (*compute the loop correction where applicable *)
										  | NONE 	 => nil)
						else
							nil) @ (computeSingleEdge t e);		
			in
				CG.addEdges(!g, List.concat (map (computeSingleEdge(sortedEdges)) sortedEdges))
			end;
		
		(* computes result graph of a single loop iteration, provided that the loop body is c and
		   that execution of all loop iterations so far have resulted with the graph g. *)
		fun doSingleIteration (g : CG.graph ref, c : CoreEx.Cmd)  =
			let
				val () = debugPrint("=======doSingleIteration=========\nIteration #"  ^ Int.toString(!itCount) ^ ":\n");
				(* Compute the graph representing the loop body, with pre-contexts which are the post contexts of the current gaph.
				   This is a graph which represents the next iteration of the loop. *)
				val g' = abstractInterpret(c, CG.postContexts(!g));
				val () = debugPrint("current graph(g):\n");
				val () = debugPrint(CG.prtEdges(!g));
				val () = debugPrint("graph to compose(g'):\n");
				val () = debugPrint(CG.prtEdges(!g'));
				(* now, compute the composition of current iteration and the next iteration graphs *)
				val g'' = !g o !g';
				val () = debugPrint("composed graph(g''):\n" ^ CG.prtEdges(g''));
				(* now - join the results of last iteration with results of current iteration *)
				val g''' = !g || g'';
				val () = debugPrint("join result graph(g'''):\n");
				val () = debugPrint(CG.prtEdges(g'''));
				val () = debugPrint("================================\n");
				val () = itCount := !itCount + 1;
				val _ = if (!debug) then TextIO.inputLine(TextIO.stdIn) else NONE;
			in
				(* return the result graph representing the next loop iteration *)
				g'''
			end;

		(* Compute the result graph of any m >0 iterations of the loop whose body is c and the iterator variable is x.
		   Update the result graph to g. This function applies loop rule L1 in the DICE paper. *)
		fun doMIterations(x : int, c : CoreEx.Cmd, g : CG.graph ref) : CG.graph ref =
			let
				(* least fixed point function *)
				fun lfp f (g : CG.graph ref, c : CoreEx.Cmd) : CG.graph ref = 
					let
						val () = debugPrint("lfp: g before:\n" ^ CG.prtEdges(!g));
						val fg = f(g, c);
						val isEq = CG.equal(fg,!g);
						val () = debugPrint(if (isEq) then "lfp reached!\n" else "lfp not reached!\ng:\n");
						val () = debugPrint(CG.prtEdges(!g));
						val () = debugPrint("fg:\n");
						val () = debugPrint(CG.prtEdges(fg));
					in
						if (isEq) then g else (g := fg; 
											   debugPrint("lfp: g after:\n" ^ CG.prtEdges(!g));
											   lfp f (g,c))
				end;
				val () = itCount := 0;
			in
				lfp doSingleIteration (g,c)
			end
			
		(* apply the loop rules. The result graph is updated in g *)
		fun doLoop(x : int, c : CoreEx.Cmd, preContexts : CG.Context list, g : CG.graph ref) : CG.graph ref =
			let
				(* compute graph representing loop's body, including a possibility of zero iterations (skip).
				   This action also applies loop rule L0 in the DICE paper 
				   (since each loop body c was expanded to "if { Skip } else { c }").*)
				val () = g := !(abstractInterpret(c, preContexts));  
				(* compute graph representing any m>0 iterations of the loop body *)
				val () = (doMIterations(x, c, g); 
				(* now, apply the loop correction operator on the result graph *)				
				          doLoopCorrection(x, g));
			in
				g
			end;
		
		(* computes the composition result graph of commands c1 and c2 and updates it to g *)
		fun doCompose(c1 : CoreEx.Cmd,c2 : CoreEx.Cmd, preContexts : CG.Context list, g : CG.graph ref) : CG.graph ref =
			let		
				val g1 = abstractInterpret(c1, preContexts);
				val g2 = abstractInterpret(c2, (CG.postContexts(!g1)));
				val () = g := (!g1 o !g2);										
			in
				g
			end;

		(* computes the join(Least Upper Bound) result graph of commands c1 and c2 and updates it to g *)			
		fun doChoice(c1 : CoreEx.Cmd,c2 : CoreEx.Cmd, preContexts : CG.Context list, g : CG.graph ref) : CG.graph ref =
			let
				val () = g := (!(abstractInterpret(c1, preContexts)) || !(abstractInterpret(c2, preContexts)))
		    in
				g
		    end
			
		(* the actual interpreter step - recursively interpret the command according to it's type and update it's attached graph with the result. *)			
		fun internalInterpret(CoreEx.cmd(c,g) : CoreEx.Cmd, preContexts : CG.Context list) : CG.graph ref =
			case c of CoreEx.Seq(c1,c2)         => doCompose(c1,c2,preContexts, g)
					|  CoreEx.Choice(c1,c2)     => doChoice(c1,c2,preContexts,g)
					|  CoreEx.Skip 		    	=> computeAssignmentGraph(1,Core.Var 1, preContexts)
					|  CoreEx.Asgn(i, e)    	=> computeAssignmentGraph(i, e, preContexts)
					|  CoreEx.Loop(i, c)    	=> doLoop(i,c,preContexts,g);

		(* returns true if command c must be re-interpreted. This will be so if current pre-contexts are 
		   not a sub-set of the pre-contexts of the graph attached to command c *)
		fun shouldUpdateCommandGraph(c,preContexts) =
			let
				val g = CoreEx.graph(c);
			in
				not (CG.preContextsIn(preContexts,!g))
			end;
		
		(* get reference to graph attached to command c *)
		val g = CoreEx.graph(c);
		val () = debugPrint("Started analysing command:\n");
		val () = debugPrint(CoreEx.prt(c));
		val () = if (shouldUpdateCommandGraph(c,preContexts)) then 
					g := !(internalInterpret(c, preContexts)) 
				 else (); (* if graph doesn't need to be updated - do nothing *)
		val () = debugPrint("Finished analysing command:\n" ^ CoreEx.prt(c) ^ "result graph:\n");
		val () = debugPrint(CG.prtEdges(!g));
		val () = debugPrint("============================================\n");
	in
		g
	end;

(* the entry point for the abstact interpreter execution. Recieves a command of the core language as input (in AST form) and 
   returns a context-graph which represents the algorithm analysis results on the input command. *)
fun analyse(c : Core.Cmd) : CG.graph = 
	let
		(* initialize "global" identifier. *)
		val () = (numVars := Core.highest_var(c));
		(* Attach a context graph to each command, and convert each loop body c to "if { Skip} else { c }". *)
		val c' = CoreEx.extendedAST(c);
	in
		(* start running the abstract interpreter with a pre-contexts list containing only the empty context, 
		   since at first we assume no contexts. *)
		!(abstractInterpret(c',[CG.makeEmptyContext(!numVars)]))
	end;
	
end; (* structure AbstractInterpreter *)



